import asyncio
import os
import re
from pathlib import Path
from typing import Optional

from browser_use import Agent, Browser
from browser_use.browser.context import BrowserContextConfig
from browser_use import BrowserConfig
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class ArticleAnalyzer:
    """文章分析器类，用于分析网页文章内容并生成结构化大纲"""
    
    def __init__(self):
        self.browser: Optional[Browser] = None
        self.llm = ChatOpenAI(
            base_url='https://api.deepseek.com/v1',
            model='deepseek-chat',
            api_key=SecretStr(os.getenv("DEEPSEEK_API_KEY")),
            temperature=0.1
        )
        
    async def setup_browser(self) -> None:
        """设置并初始化浏览器"""
        print("正在初始化浏览器配置...")
        
        # 配置浏览器上下文
        context_config = BrowserContextConfig(
            # 页面加载设置
            minimum_wait_page_load_time=0.5,  # 最小等待时间
            wait_for_network_idle_page_load_time=3.0,  # 等待网络空闲时间
            maximum_wait_page_load_time=30.0,  # 最大等待时间
            
            # 显示设置
            browser_window_size={'width': 1280, 'height': 1100},  # 浏览器窗口大小
            locale='zh-CN',  # 设置中文环境
            highlight_elements=True,  # 高亮显示可交互元素
            viewport_expansion=800,  # 视口扩展
            
            # 用户代理设置
            user_agent='Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36'
        )
        
        # 浏览器基础配置
        browser_config = BrowserConfig(
            headless=False,  # 显示浏览器界面
            disable_security=True,  # 禁用安全限制
            extra_chromium_args=[
                '--no-sandbox',
                '--disable-dev-shm-usage',
                '--disable-blink-features=AutomationControlled',
                '--start-maximized',
                '--disable-web-security',  # 禁用网页安全策略
                '--disable-features=IsolateOrigins,site-per-process'  # 禁用站点隔离
            ],
            new_context_config=context_config
        )
        
        try:
            print("正在启动浏览器...")
            self.browser = Browser(config=browser_config)
            self.context = await self.browser.new_context()
            print("浏览器初始化成功！")
        except Exception as e:
            raise RuntimeError(f"浏览器初始化失败: {str(e)}")

    def clean_markdown_content(self, content: str) -> str:
        """清理和格式化 markdown 内容"""
        if not content:
            print("警告：传入的内容为空")
            return ""
            
        try:
            # 如果内容是字符串形式的字典，尝试提取 extracted_content
            if isinstance(content, str) and content.strip().startswith('{'):
                try:
                    import json
                    print("尝试解析JSON内容...")
                    data = json.loads(content)
                    if isinstance(data, dict):
                        if 'extracted_content' in data:
                            content = data['extracted_content']
                            print("从JSON中提取了extracted_content")
                        elif 'content' in data:
                            content = data['content']
                            print("从JSON中提取了content")
                except json.JSONDecodeError as e:
                    print(f"JSON解析失败: {str(e)}")
                    pass

            # 确保内容是字符串
            content = str(content)
            print(f"开始清理的内容长度: {len(content)}")
            
            # 如果内容中没有标题标记，尝试查找可能的标题
            if not re.search(r'^#', content, re.MULTILINE):
                print("未找到标题标记，尝试提取标题...")
                # 查找可能的标题（第一个非空行）
                first_line = next((line for line in content.split('\n') if line.strip()), '')
                if first_line:
                    content = f"# {first_line.strip()}\n\n{content}"
                    print("已添加标题标记")
            
            # 清理多余的空行和格式化内容
            lines = []
            in_code_block = False
            code_block_indent = ""
            
            for line in content.split('\n'):
                # 处理代码块
                if line.strip().startswith('```'):
                    in_code_block = not in_code_block
                    if in_code_block:
                        code_block_indent = re.match(r'^\s*', line).group()
                    lines.append(line)
                    continue
                    
                # 在代码块内的内容保持原样
                if in_code_block:
                    if line.startswith(code_block_indent):
                        line = line[len(code_block_indent):]
                    lines.append(code_block_indent + line)
                    continue
                
                line = line.strip()
                
                # 跳过不需要的内容
                if any(skip in line for skip in ['AgentHistoryList', 'extracted_content']):
                    continue
                    
                # 处理标题行
                if line.startswith('#'):
                    if lines and not lines[-1].strip():
                        lines.append(line)
                    else:
                        lines.extend(['', line])
                    lines.append('')
                # 处理列表项
                elif line.startswith(('-', '*', '+')):
                    lines.append(line)
                # 处理引用
                elif line.startswith('>'):
                    if lines and not lines[-1].startswith('>'):
                        lines.append('')
                    lines.append(line)
                    if len(lines) > 1 and not lines[-2].startswith('>'):
                        lines.append('')
                # 处理其他非空行
                elif line:
                    lines.append(line)
                # 保留必要的空行
                elif lines and lines[-1].strip():
                    lines.append('')
            
            # 合并行
            content = '\n'.join(lines)
            
            # 清理连续的空行
            content = re.sub(r'\n{3,}', '\n\n', content)
            
            # 确保内容以换行符结束
            if not content.endswith('\n'):
                content += '\n'
            
            # 清理开头和结尾的空行
            content = content.strip() + '\n'
            
            print(f"清理后的内容长度: {len(content)}")
            
            # 如果清理后的内容为空，返回提示信息
            if not content.strip():
                return "# 无法提取有效内容\n\n可能的原因：\n- 页面加载不完整\n- 内容被保护\n- 格式不支持\n"
            
            return content
            
        except Exception as e:
            print(f"清理内容时出错: {str(e)}")
            # 如果处理失败，返回原始内容
            return str(content)

    async def analyze_url(self, url: str, timeout: int = 180) -> Optional[str]:
        """分析指定URL的文章内容"""
        if not url.startswith(('http://', 'https://')):
            url = 'https://' + url
            
        print(f"正在访问URL: {url}")
        
        try:
            if not self.browser:
                await self.setup_browser()
                
            agent = Agent(
                task=(
                    f'任务：提取并分析文章内容\n'
                    f'步骤：\n'
                    f'1. 访问URL：{url}\n'
                    f'2. 等待页面加载完成\n'
                    f'3. 定位文章主体内容：\n'
                    f'   - 查找文章标题元素\n'
                    f'   - 查找文章正文容器\n'
                    f'   - 识别文章大纲或目录结构\n'
                    f'4. 提取内容并转换为Markdown：\n'
                    f'   - 标题使用 # 格式\n'
                    f'   - 章节使用 ## 和 ### 格式\n'
                    f'   - 重点内容使用 - 列表\n'
                    f'   - 保留原文代码块和引用\n'
                    f'   - 保持段落结构完整\n'
                    f'5. 清理无关内容：\n'
                    f'   - 移除广告\n'
                    f'   - 移除评论区\n'
                    f'   - 移除导航栏\n'
                    f'   - 移除页脚\n'
                    f'6. 返回处理后的Markdown文本\n'
                    f'要求：\n'
                    f'- 确保提取完整的文章内容\n'
                    f'- 保持文章结构清晰\n'
                    f'- 内容必须有实际文本\n'
                    f'- 直接返回Markdown格式文本'
                ),
                llm=self.llm,
                browser=self.browser,
                browser_context=self.context,
                use_vision=False
            )
            
            try:
                print("正在提取文章内容...")
                result = await asyncio.wait_for(agent.run(), timeout=timeout)
                print("内容提取完成，正在处理...")
            except asyncio.TimeoutError:
                print(f"提取超时（{timeout}秒），正在重试...")
                # 如果超时，尝试获取已有的内容
                result = agent.get_current_result()
                if not result:
                    print("未能获取任何内容")
                    return None
            
            # 获取并处理内容
            content = None
            
            # 打印调试信息
            print(f"结果类型: {type(result)}")
            
            if isinstance(result, dict):
                print("处理字典类型结果...")
                if 'extracted_content' in result:
                    content = result['extracted_content']
                elif 'content' in result:
                    content = result['content']
                else:
                    content = str(result)
                print(f"从字典中提取的内容长度: {len(str(content)) if content else 0}")
            elif hasattr(result, 'extracted_content'):
                print("处理带extracted_content属性的结果...")
                content = result.extracted_content
            elif hasattr(result, 'content'):
                print("处理带content属性的结果...")
                content = result.content
            elif callable(result):
                print("处理可调用结果...")
                try:
                    content = result()
                except Exception as e:
                    print(f"调用结果函数时出错: {str(e)}")
                    content = str(result)
            else:
                print("处理其他类型结果...")
                content = str(result)
            
            if not content:
                print("警告：提取的内容为空")
                return None
                
            print(f"提取的原始内容长度: {len(str(content))}")
            
            try:
                # 确保 content 是字符串
                if not isinstance(content, str):
                    content = str(content)
                
                # 检查内容是否为空
                if not content.strip():
                    print("警告：提取的内容为空字符串")
                    return None
                
                # 清理和格式化内容
                cleaned_content = self.clean_markdown_content(content)
                print(f"清理后的内容长度: {len(cleaned_content)}")
                
                if not cleaned_content.strip():
                    print("警告：清理后的内容为空")
                    return None
                
                # 保存结果到文件
                output_path = Path("article_structure.md")
                try:
                    # 确保内容是UTF-8编码
                    output_path.write_text(cleaned_content, encoding="utf-8")
                    print(f"\n分析完成！结果已保存到：{output_path}")
                    print("\n生成的文章结构如下：\n")
                    print(cleaned_content)
                    return cleaned_content
                except Exception as e:
                    print(f"保存文件时出错: {str(e)}")
                    return cleaned_content
            except Exception as e:
                print(f"处理内容时出错: {str(e)}")
                if content and isinstance(content, str):
                    return content
                return str(content)
            
        except Exception as e:
            print(f"分析过程发生错误: {str(e)}")
            return None
        finally:
            try:
                # 每次分析完成后关闭浏览器，避免资源占用
                if self.browser:
                    await self.context.close()
                    await self.browser.close()
                    self.browser = None
                    print("浏览器已关闭")
            except Exception as e:
                print(f"关闭浏览器时出错: {str(e)}")

async def main():
    """主程序入口"""
    if not os.getenv("DEEPSEEK_API_KEY"):
        print("错误：需要设置DEEPSEEK_API_KEY环境变量")
        return
        
    analyzer = ArticleAnalyzer()
    
    try:
        while True:
            url = input("\n请输入要分析的文章URL (输入q退出): ").strip()
            
            if url.lower() == 'q':
                break
                
            if not url:
                print("错误：URL不能为空")
                continue
                
            result = await analyzer.analyze_url(url)
            if result:
                print("\n生成的文章结构如下：\n")
                print(result)
                
            retry = input("\n是否继续分析其他文章？(y/n): ").lower()
            if retry != 'y':
                break
                
    finally:
        await analyzer.close()

if __name__ == "__main__":
    asyncio.run(main()) 